from hmmlearn import hmm
import numpy as np
states = ["box1","box2","box3"]
n_states = len(states)
observations = ["red","white"]
n_observations = len(observations)
start_probability = np.array([0.2,0.4,0.4])
transition_probability = np.array([
    [0.5,0.2,0.3],
    [0.3,0.5,0.2],
    [0.2,0.3,0.5]
])
emission_probability = np.array([
    [0.5,0.5],
    [0.4,0.6],
    [0.7,0.3]
])

model = hmm.CategoricalHMM(n_components=n_states)
model.startprob_ = start_probability
model.transmat_ = transition_probability
model.emissionprob_ = emission_probability




seen =np.array([[1,1,0]]).T
logprob,box = model.decode(seen,algorithm="viterbi")
print(np.array(states)[box])
box2 = model.predict(seen)
print(np.array(states)[box2])
print(model.score(seen))

model2 = hmm. CategoricalHMM(n_components=n_states,n_iter=20,tol=0.01)
X2 =np. array([[0, 1, 0,1],[0, 0, 0, 1],[1,0,1,1]])
model2.fit(X2)
print('startprob_' , model2.startprob_)
print('----')
print('transmat_' , model2.transmat_)
print('------------')
print('emissionprob_' , model2.emissionprob_)
print('---------------------------')
print('score' , model2.score(X2))